import os
import random
from PIL import Image

# 配置参数
dataset_path = r"F:\人工智能教材编写\traffic_sign\train"
output_filename = "dataset_overview.png"
grid_rows = 4    # 网格行数
grid_cols = 5     # 网格列数
tile_size = (128, 128)  # 单张示例图尺寸

# 收集所有图片路径
all_images = []
for class_dir in os.listdir(dataset_path):
    class_path = os.path.join(dataset_path, class_dir)
    if os.path.isdir(class_path):
        images = [os.path.join(class_path, f)
                 for f in os.listdir(class_path)
                 if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        all_images.extend(images)

# 随机采样20张图片
selected = random.sample(all_images, grid_rows * grid_cols)

canvas = Image.new('RGB',
                 (tile_size[0] * grid_cols,  # 宽度 = 128 * 5
                  tile_size[1] * grid_rows)) # 高度 = 128 * 4

# 处理并排列图片
for index, img_path in enumerate(selected):
    img = Image.open(img_path).convert('RGB')
    img = img.resize(tile_size)

    # 计算放置位置
    row = index // grid_cols
    col = index % grid_cols
    canvas.paste(img, (col * tile_size[0], row * tile_size[1]))

# 添加黑色边框
canvas_with_border = Image.new('RGB',
                              (canvas.width + 2, canvas.height + 2),
                              color=(0, 0, 0))
canvas_with_border.paste(canvas, (1, 1))

# 保存结果
canvas_with_border.save(output_filename)
print(f"数据集示例图已生成：{output_filename}")
